#ifndef AVL_hpp
#define AVL_hpp
#endif /* AVL_hpp */

#include<iostream>
#include<cassert>
using namespace std;

template<class K,class V>
struct AVLTreeNode
{
    AVLTreeNode<K,V>* _left;
    AVLTreeNode<K,V>* _right;
    AVLTreeNode<K,V>* _parent;
    
    pair<K,V> _kv;
    int _bf; //平衡因子
    
    AVLTreeNode(const pair<K, V>&kv)
    :_left(nullptr)
    ,_right(nullptr)
    ,_parent(nullptr)
    ,_kv(kv)
    ,_bf(0)
    {}
};

template<class K,class V>
class AVLTree
{
    typedef AVLTreeNode<K,V> Node;
public:
    bool Insert(const pair<K, V>& kv)
    {
        if(_root==nullptr)
        {
            _root=new Node(kv);
            return true;
        }
        
        Node* parent=nullptr;
        Node* cur=_root;
        
        while(cur)
        {
            if(cur->_kv.first<kv.first)
            {
                parent=cur;
                cur=cur->_right;
            }
            else if(cur->_kv.first>kv.first)
            {
                parent=cur;
                cur=cur->_left;
            }
            else
            {
                return false;
            }
        }
        cur=new Node(kv);
        if(parent->_kv.first<kv.first)
        {
            parent->_right=cur;
            cur->_parent=parent;//将节点的父指针指向父亲
        }
        else
        {
            parent->_left=cur;
            cur->_parent=parent;//将节点的父指针指向父亲
        }
        
        
        //这里之前，和以前二叉搜索树步骤差不多，
        //这里之后，就要实现平衡的问题
        while(parent)//parent不能为空
        {
            //更新父节点
            //如果是父节点的左节点，则父节点的平衡因子--
            if(cur==parent->_left)
                parent->_bf--;
            //如果是父节点的右节点，则父节点的平衡因子++
            else
                parent->_bf++;
            
            //更新完父节点的平衡因子后，需要判断是否继续向上更新
            
            //1.如果父节点的平衡因子为0，则不需要向上更新
            if(parent->_bf==0)
                break;
            //2.如果父节点的平衡因子为-1 or 1，则需要向上更新
            else if(parent->_bf==1||parent->_bf==-1)
            {
                cur=parent;//子节点指向父节点
                parent=cur->_parent;//父节点指向祖先节点
            }
            
            //3.如果父节点的平衡因子为2 or -2 则需要调整整个搜索二叉树
            else if(parent->_bf==2||parent->_bf==-2)
            {
                //旋转
                
                //
                if(parent->_bf==2&&cur->_bf==1)
                {
                    RotateL(parent);
                }
                
                else if(parent->_bf==-2&&cur->_bf==-1)
                {
                    RotateR(parent);
                }
                
                else if(parent->_bf==2&&cur->_bf==-1)
                {
                    RotateRL(parent);
                }
                
                else if(parent->_bf==-2&&cur->_bf==1)
                {
                    RotateLR(parent);
                }
                
                //1）旋转让这颗子树平衡
                //2）旋转本质也降低了高度
                break;
            }
            //平衡因子不可能为大于2或小于-2
            else
            {
                assert(false);
            }
        }
        
        
        return true;
    }
    
    void RotateL(Node* parent)//左单旋
    {
        Node* subR=parent->_right;
        Node* subRL=subR->_left;
        //用于记录parent的parent，因为左旋转的时候，不一定parent是根
        Node* parentParent=parent->_parent;
        
        parent->_right=subRL;
        subR->_left=parent;
        
        parent->_parent=subR;
        if(subRL)
            subRL->_parent=parent;
        
        //如果parent是整个树的根
        if(_root==parent)
        {
            _root=subR;
            subR->_parent=nullptr;
        }
        //如果parent不是是整个树的根
        else
        {
            //将parentParent的左指向subR
            if(parentParent->_left==parent)
                parentParent->_left=subR;
            //将parentParent的右指向subR
            else
                parentParent->_right=subR;
            //subR的parent指向parentParent
            subR->_parent=parentParent;
        }
        parent->_bf=0;//parent平衡因子为0
        subR->_bf=0;//subR平衡因子为0
    }
    
    void RotateR(Node* parent)//右单旋
    {
        Node* subL=parent->_left;
        Node* subLR=subL->_right;
        Node* parentParent=parent->_parent;
        
        parent->_left=subLR;
        subL->_right=parent;
        
        parent->_parent=subL;
        if(subLR)
            subLR->_parent=parent;
        
        if(_root==parent)
        {
            _root=subL;
            subL->_parent=nullptr;
        }
        else
        {
            if(parentParent->_left==parent)
                parentParent->_left=subL;
            else
                parentParent->_right=subL;
            subL->_parent=parentParent;
        }
        
        parent->_bf=0;
        subL->_bf=0;
    }
    void RotateRL(Node* parent)//右左双旋
    {
        Node* subR=parent->_right;
        Node* subRL=subR->_left;
        
        int bf=subRL->_bf;
        
        RotateR(parent->_right);
        RotateL(parent);
        
        if(bf==0)
        {
            //subRL自己就是新增
            parent->_bf=subR->_bf=subRL->_bf=0;
        }
        else if(bf==-1)
        {
            //subRL左子树新增
            parent->_bf=0;
            subRL->_bf=0;
            subR->_bf=1;
        }
        else if(bf==1)
        {
            //subRL右子树新增
            parent->_bf=-1;
            subRL->_bf=0;
            subR->_bf=0;
        }
        else
        {
            assert(false);
        }
    }
    
    void RotateLR(Node* parent)//右左双旋
    {
        Node* subL=parent->_left;
        Node* subLR=subL->_right;
        int bf=subLR->_bf;//记录subLR的bf值，然后判断他是那种情况
        
        RotateL(parent->_left);
        RotateR(parent);
        
        if(bf==0)
        {
            subL->_bf=subLR->_bf=parent->_bf;
        }
        
        else if(bf==1)
        {
            //subLR右子树新增
            parent->_bf=0;
            subLR->_bf=0;
            subL->_bf=-1;
        }
        else if(bf==-1)
        {
            //subLR左子树新增
            parent->_bf=1;
            subLR->_bf=0;
            subL->_bf=0;
        }
        else
        {
            assert(false);
        }
    }
    
    //中序遍历
    void _InOder(Node* root)
    {
        if(root==nullptr)
            return;
        _InOder(root->_left);
        cout<<root->_kv.first<<" ";
        _InOder(root->_right);
    }
    void InOrder()
    {
        _InOder(_root);
        cout<<endl;
    }
    
    size_t _Size(Node* root)
    {
        if(root==nullptr)
            return 0;
        return _Size(root->_left)+_Size(root->_right)+1;
    }
    
    size_t Size()
    {
        return  _Size(_root);
    }
    
    bool Find(const K& key)
    {
        Node* cur=_root;
        while(cur)
        {
            if(cur->_kv.first<key)
                cur=cur->_right;
            else if(cur->_kv.first>key)
                cur=cur->_left;
            else
                return true;
        }
        return  false;
    }
private:
    Node* _root=nullptr;
};
